import os
from ai_butler_sdk.celery_app import celery_app
from ai_butler_sdk.train import TrainBase
from trainer import Trainer

from ai_butler_sdk.utils import xml_to_txt


class MockTrain(TrainBase):

    def pre_train(self):
        self.txt_dataset_dir = os.path.join(self.root_path, 'new_datasets')
        os.makedirs(self.txt_dataset_dir, exist_ok=True)
        xml_to_txt(self.data_sets_local_path, self.txt_dataset_dir)

        self.coco_dir = os.path.join(self.root_path, f"coco_{self.train_task_id}")
        os.makedirs(self.coco_dir, exist_ok=True)

    def train(self):
        trainer = Trainer(
            task_id=self.train_task_id,
            raw_dataset=self.txt_dataset_dir,
            coco_dir=self.coco_dir,
            output_dir=self.root_path,
            train_params=self.train_params,
            model_name=self.network,
            second_train=bool(self.pretrain_model_weight_download_url),  # 是否有
            train_result_zip=self.result_local_path,  # result.zip
            train_log=self.log_local_path  # train.log
        )

        trainer.train()
        trainer.delete()


# @celery_app.task
# def train(
#         train_task_id: str,
#         data_set_urls: list[str],
#         train_params: dict,
#         log_upload_url: str,
#         model_weight_upload_url: str,
#         pretrain_model_weight_download_url: str | None = None,
#
# ):
#     """模拟训练"""
#     mock_train = MockTrain(
#         train_task_id,
#         data_set_urls,
#         train_params,
#         log_upload_url,
#         model_weight_upload_url,
#         pretrain_model_weight_download_url,
#     )
#     mock_train()




if __name__ == '__main__':
    data = {'train_task_id': '5', 'network': 'yolov5n', 'data_set_urls': ['http://test-ai.shuanzhineng.com:9000/aibutler/admin/datasets/2024-03/sada.zip?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIOSFODNN7EXAMPLE%2F20240319%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240319T100026Z&X-Amz-Expires=604800&X-Amz-SignedHeaders=host&X-Amz-Signature=87713a52bfc993e569195e9d1e1900fafea1ea598ae9ee425dd36137b8327a0b'], 'pretrain_model_weight_download_url': None, 'train_params': {'train_data_ratio': 0.8, 'epochs': 1, 'batch_size': 16, 'imgsz': 224, 'save_period': -1, 'seed': 0, 'device': '0', 'multi_scale': False, 'workers': 0, 'cos_lr': False, 'label_smoothing': 0.0, 'freeze': [0], 'optimizer': 'SGD', 'train_hyp_params': {}}, 'model_weight_upload_url': 'http://test-ai.shuanzhineng.com:9000/aibutler/admin/train/2024-03/5/result.zip?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIOSFODNN7EXAMPLE%2F20240315%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240315T091453Z&X-Amz-Expires=604800&X-Amz-SignedHeaders=host&X-Amz-Signature=737b5e1784ad180f5404c82b0b3ce48163fbc6c90591023287f28afad22f02a5', 'log_upload_url': 'http://test-ai.shuanzhineng.com:9000/aibutler/admin/train/2024-03/5/train.log?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIOSFODNN7EXAMPLE%2F20240315%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240315T091453Z&X-Amz-Expires=604800&X-Amz-SignedHeaders=host&X-Amz-Signature=79feb7e595e2a0ff1cc2ba1dc3fa5862f8fea37e38e222a476e5333a292347cd'}
    train = MockTrain(**data)
    train()
